from langchain.document_loaders import JSONLoader
from faissManager import FAISS
from sent2vec import Sent2VecEmbeddings
from argparse import ArgumentParser
from accelerate import Accelerator

def parse_args():
    parser=ArgumentParser()
    parser.add_argument('--input',type=str,help='指定需要处理的json文件',default='/home/lxy/wxbdata/db/merge.json')
    parser.add_argument('--save_dir',type=str,help='指定索引和key被转成句向量的张量保存路径',default='/home/lxy/multiR/gte_11m21d')
    parser.add_argument('--kv',type=bool,help='用于确定数据库索引的对象是相同键值还是不同键值',default=True,)
    parser.add_argument('--k_name',type=str,help='指定json中哪个键作为key',default='question')
    parser.add_argument('--v_name',type=str,help='指定json中哪个键作为value',default='answer')
    parser.add_argument('--model_name',type=str,help='指定句向量模型,需要是hf格式',default='/data/lxy/thenlper/gte-large-zh')
    parser.add_argument('--batch_size',type=int,help='编码句向量的batch_size',default=128)
    parser.add_argument('--reuse',type=bool,help='是否重用save_dir保存的张量，不需要重复计算',default=False)
    return parser.parse_args()

def MetaDataExtractor(sample, metadata):
    # 通常情况下这个函数并不会对q文档使用，因为他的用途只是被编码成向量形式
    if sample:
        if metadata is not None:
          metadata['human']=sample.get('human',0)
          metadata['question']=sample.get('question',None)
        else:
          metadata={'human':sample.get('human',0)}
          metadata={'question':sample.get('question',None)}

    return metadata

  

if __name__ == '__main__':
    args=parse_args()

    q_documents = JSONLoader(args.input,jq_schema='.[]',content_key=args.k_name).load()
    print('q_documents load success')

    a_documents = JSONLoader(args.input,jq_schema='.[]',content_key=args.v_name,metadata_func=MetaDataExtractor).load()
    print('a_documents load success')

    accelerator = Accelerator()
    with accelerator.main_process_first():
      embedding=Sent2VecEmbeddings(model_name=args.model_name,batch_size=args.batch_size)

    db = FAISS.from_idxType_and_documents(documents=a_documents,
                                          embedding= embedding,
                                          tensorSaveDir=args.save_dir,qa=args.kv,query_document=q_documents,reuse=False)
                                        #   index_type='Flat',do_max_norm=False,normalize_L2=True)


    db.save_local(args.save_dir)